from rkllm import *
import ctypes
import signal
from threading import Thread
import multiprocessing as mp
from ctypes import c_void_p, c_char_p, c_int32, c_float, POINTER, memset, byref
from rknnlite.api import RKNNLite
import cv2
import numpy as np

PROMPT_TEXT_PREFIX = "<用户>"
PROMPT_TEXT_POSTFIX = "\n<AI>"


class MiniCPMV(object):
    rknn_lite = RKNNLite()  
    llmHandle = RKLLM_Handle_t()

    def __init__(self, model_path, vision_path):
        self.callback = callback_type(self.result_callback)
        self.init_rkllm(model_path)
        self.rkllm_infer_params = self.get_infer_params()
        self.init_rknn(vision_path)

        self.full_output = ""


    def result_callback(self, result, userdata, state):
        if state == LLMCallState.RKLLM_RUN_FINISH:
            print()
        elif state == LLMCallState.RKLLM_RUN_ERROR:
            print("run error")
        elif state == LLMCallState.RKLLM_RUN_GET_LAST_HIDDEN_LAYER:
            hidden_layer = result.contents.last_hidden_layer
            if hidden_layer.embd_size and hidden_layer.num_tokens:
                data_size = hidden_layer.embd_size * hidden_layer.num_tokens * ctypes.sizeof(c_float)
                print(f"data_size:{data_size}")
                with open("last_hidden_layer.bin", "wb") as f:
                    f.write(ctypes.string_at(hidden_layer.hidden_states, data_size))
                print("Data saved to output.bin successfully!")
        elif state == LLMCallState.RKLLM_RUN_NORMAL:
            output_text = result.contents.text.decode('utf-8')
            self.full_output += output_text
            print(output_text, end="", flush=True)


    def init_rkllm(self, model_path):
        rkllm_param = RKLLMParam()
        rkllm_param.model_path = bytes(model_path, 'utf-8')

        rkllm_param.max_context_len = 1024
        rkllm_param.max_new_tokens = 512
        rkllm_param.skip_special_token = True

        rkllm_param.top_k = 30
        rkllm_param.top_p = 0.95
        rkllm_param.temperature = 1.0
        rkllm_param.repeat_penalty = 1.1
        rkllm_param.frequency_penalty = 0.0
        rkllm_param.presence_penalty = 0.0

        rkllm_param.img_start = b"<image>"
        rkllm_param.img_end = "</image>\n".encode('utf-8')
        rkllm_param.img_content = "<unk>".encode('utf-8')
        rkllm_lib.rkllm_init(ctypes.byref(self.llmHandle), ctypes.byref(rkllm_param), self.callback)
        

    def init_rknn(self, vision_path):
        ret = self.rknn_lite.load_rknn(vision_path)
        if ret != 0:
            print('Load RKNN model failed')
            exit(ret)
        print('===>: load rknn done!')
        ret = self.rknn_lite.init_runtime()
        if ret != 0:
            print('Init runtime failed')
            exit(ret)
        print('===>: init runtime done!')

    def get_infer_params(self):
        rkllm_infer_params = RKLLMInferParam()
        ctypes.memset(ctypes.byref(rkllm_infer_params), 0, ctypes.sizeof(RKLLMInferParam))
        rkllm_infer_params.mode = RKLLMInferMode.RKLLM_INFER_GENERATE
        return rkllm_infer_params

    def prepare_image(self, img):
        h, w, _ = img.shape
        scale = 384 / max(h, w)
        new_w, new_h = int(w * scale), int(h * scale)
        img_resized = cv2.resize(img, (new_w, new_h))

        # 创建黑色背景并将缩放图像居中粘贴
        img_centered = np.zeros((384, 384, 3), dtype=np.uint8)
        x_offset = (384 - new_w) // 2
        y_offset = (384 - new_h) // 2
        img_centered[y_offset:y_offset+new_h, x_offset:x_offset+new_w, :] = img_resized

        # 转换为 RGB 通道顺序
        img_centered = cv2.cvtColor(img_centered, cv2.COLOR_BGR2RGB)
        img_centered = np.expand_dims(img_centered, axis=0)
        return img_centered

    def run(self,img, prompt):
        img_embedding = self.rknn_lite.inference(inputs=[img])[0].astype(np.float32)
        input_str = PROMPT_TEXT_PREFIX + prompt + PROMPT_TEXT_POSTFIX
        rkllm_input = RKLLMInput()
        rkllm_input.input_type = 3  # RKLLM_INPUT_MULTIMODAL
        rkllm_input.input_data.multimodal_input.prompt = ctypes.c_char_p(input_str.encode('utf-8'))
        rkllm_input.input_data.multimodal_input.image_embed = img_embedding.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
        rkllm_input.input_data.multimodal_input.n_image_tokens = 64

        self.full_output = ""
        rkllm_lib.rkllm_run(self.llmHandle, ctypes.byref(rkllm_input), ctypes.byref(self.rkllm_infer_params), None)

    def release(self):
        rkllm_lib.rkllm_destroy(self.llmHandle)

